from torch import nn
import cv2
import torch.nn.functional as F
import numpy as np
import torch
import math

class Features_SIFT(nn.Module):
    def __init__(self, *args):
        super().__init__()
        self.getter = cv2.SIFT_create(1000)

    # def forward(self, x):
    #     # getter = cv2.SIFT_create()
    #     x = x.numpy() * 255
    #     kps_left, des_left = self.getter.detectAndCompute(x,None)
    #     return {"keypoints":kps_left,"descriptors":des_left}
    def forward(self, input):
        # getter = cv2.SIFT_create()

        image = input['image'][0].cpu().numpy()
        # image = cv2.imread(imgpath)


        kps_left, des_left = self.getter.detectAndCompute(image,None)
        
        kps_arr = np.zeros((len(kps_left),2),dtype=np.float32)
        for ind in range(len(kps_left)):
            kps_arr[ind][0] = kps_left[ind].pt[0]
            kps_arr[ind][1] = kps_left[ind].pt[1]


        return {"keypoints":kps_arr,"descriptors":des_left}

class Features_SIFT_WiITH_MASK(nn.Module):
    def __init__(self, *args):
        super().__init__()
        self.getter = cv2.SIFT_create(1000)

    def forward(self, input):
        # getter = cv2.SIFT_create()

        imgpath = input['image_path'][0]
        maskpath = input['mask_path'][0]
        image = cv2.imread(imgpath)
        mask =  cv2.imread(maskpath).astype(np.float32)


        kps_left, des_left = self.getter.detectAndCompute(image,None)
        
        mask = torch.tensor(mask)
        mask2 = F.max_pool2d(mask*-1,kernel_size=5,stride=1,padding=2)
        mask2 *= -1
        mask_nparr = mask2.numpy()


        kps_arr = np.zeros((len(kps_left),2),dtype=np.float32)
        for ind in range(len(kps_left)):
            kps_arr[ind][0] = kps_left[ind].pt[0]
            kps_arr[ind][1] = kps_left[ind].pt[1]

        points_cnt = kps_arr.shape[0]

        id_list = []
        for i in range(points_cnt):
            pt = kps_arr[i]
            col = int(math.floor(pt[0]))#x
            row = int(math.floor(pt[1]))#y

            #行在前，列在后
            mv = mask_nparr[row,col]
            if mv[0] > 0:
                id_list.append(i)
        if len(id_list) > 0:
            keypoints2 = kps_arr[id_list,:]
            descriptors2 = des_left[id_list,:]
            return {"keypoints":keypoints2,"descriptors":descriptors2}
        else:
            return {"keypoints":None,"descriptors":None}